import torch
from torch import nn

__all__ = ['deit']

from timm.models.vision_transformer import vit_deit_base_patch16_224, vit_base_patch16_224, vit_deit_tiny_patch16_224, \
    vit_deit_small_patch16_224, vit_deit_base_patch16_224, vit_base_resnet50_224_in21k
from timm.models.resnet import resnet18, resnet50


class DeiT(nn.Module):


    def __init__(self, pretrained, in_chans, num_classes=2, num_features=2):
        super(DeiT, self).__init__()

        self.num_features = num_features
        self.num_classes = num_classes

        self.rgb_stream = vit_base_resnet50_224_in21k(num_classes=self.num_features, pretrained=True)
        # self.rgb_stream = resnet18(num_classes=self.num_features, pretrained=True)
        # self.msr_stream = vit_deit_base_patch16_224(num_classes=self.num_features, pretrained=True)

        # self.classifier = nn.Linear(self.num_features, self.num_classes)

    def forward(self, x_rgb, x_msr):
        # x_rgb = self.rgb_stream(x_rgb)
        x_rgb = self.rgb_stream(x_rgb, x_msr)
        # x_rgb = self.rgb_stream(x_rgb)
        print("x_rgb.shape: ", x_rgb.shape, "x_rgb: ", x_rgb)

        # x_msr = self.msr_stream(x_msr)

        # x = self.classifier(x_rgb)
        # print("x.shape: ", x.shape, "x: ", x)

        # output = []
        #
        # output.append(x_rgb)
        # output.append(x)
        # for i, feat in enumerate(output):
        #     print("第{}个:".format(i+1), feat.shape)
        #     print("第{}个:".format(i+1), feat)
        return x_rgb


def deit(in_chans=3, pretrained=False):
    model = DeiT(in_chans=in_chans, pretrained=pretrained)
    model.default_cfg = model.rgb_stream.default_cfg
    return model


if __name__ == "__main__":
    model = DeiT(in_chans=3, num_classes=2, pretrained=True)
    x1 = torch.randn(8, 3, 224, 224)
    x2 = torch.randn(8, 3, 224, 224)
    regression = model(x1, x2)